import numpy as np

def s_box_substitution(data, key_stream):
    # 创建32x32的S盒
    s_box = np.array([[0x6, 0x3, 0xa, 0x5, 0xc, 0x8, 0x1, 0xb],
                      [0x9, 0xe, 0x2, 0x7, 0x0, 0x4, 0xd, 0xf]])
    
    # 生成密钥流
    key_stream = np.array(key_stream, dtype=np.uint8)
    
    # 执行S盒置换
    data = s_box[data // 8, data % 8]
    
    # 执行ARX操作
    data = np.bitwise_xor(data, key_stream)
    data = np.bitwise_and(data, 0xFF)
    data = np.bitwise_xor(data, np.roll(data, 1, axis=1))
    data = np.bitwise_xor(data, np.roll(data, 1, axis=0))
    
    return data

# 测试
plaintext = "Hello World!"
key_stream = [0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0]

# 将明文转换为32x32的整数数组
data = np.frombuffer(plaintext.encode(), dtype=np.uint8)
data = np.pad(data, (0, 32*32 - len(data)), mode='constant')
data = data.reshape(32, 32)

# 执行S盒置换和ARX操作
encrypted_data = s_box_substitution(data, key_stream)

# 输出加密后的结果
print("Encrypted Data:")
print(encrypted_data)
